import sys 
sys.path.append('code')
from feature_eng.alpha_series.alpha_ser import Alpha_ser0, Alpha_ser1, Alpha_ser2, Alpha_ser3, Alpha_ser_stastic
from feature_eng.alpha_series.alpha_df import Alpha_df0, Alpha_df1, Alpha_df2, Alpha_df3, Alpha_df_stat
import os  
from joblib import Parallel, delayed
import itertools
import pandas as pd
import time

class Ser_main():
    
    def __init__(self, columns_root=None, save_root=None, price_root=None, value_root=None) -> None:
        self.price_root = r'data/stock_data/daily/price' if price_root is None else  price_root
        self.value_root = r'data/stock_data/daily/market_value' if value_root is None else  value_root
        self.columns_root = r'data/stock_data/daily/consentrate_price/close.pkl.gzip' if columns_root is None else columns_root        
        self.save_root = r'data/cmodty/my_feature/feature_data_ser' if save_root is None else save_root     
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)
        self.column = pd.read_pickle(self.columns_root).columns
            
    def calc_alpha(self, func, name='alpha_ser_0', *params):
        try :
            t1 = time.time()
            lst = []
            for code in self.column:
                price_file = os.path.join(self.price_root,f'{code}.pkl.gzip')
                value_file = os.path.join(self.value_root,f'{code}.pkl.gzip')
                price = pd.read_pickle(price_file)
                value = pd.read_pickle(value_file)
                df_i = pd.concat([price,value], axis=1)
                df_i = df_i
                alpha_i = pd.DataFrame(index = df_i.index) 
                ary = func(name, df_i, params)
                alpha_i[code] = ary
                lst.append(alpha_i)
            df = pd.concat(lst, axis=1)
            df.columns = self.column
            df = df[df.index>'2005-01-01']
            save_file = os.path.join(self.save_root, f'{name}_{params}.pkl.gzip')
            df.to_pickle(save_file)
            print(f"{name}_{params} calculating cost time:{(time.time()-t1):.2f}s")
        except:
            with open(r"code\feature_eng\alpha_series\info.txt", 'a', encoding="utf-8") as f:
                f.write(f"\n{name}_{params}")
            print(f"{name}_{params}_{params} calculating faild")

    def eval_alpha_ser0(self, name, df_i, params):
        alp1 = Alpha_ser0(df_i)
        ary = eval(f'alp1.{name}')()
        return ary

    def eval_alpha_ser1(self, name, df_i, params):
        alp1 = Alpha_ser1(df_i)
        ary = eval(f'alp1.{name}')(params[0])
        return ary

    def eval_alpha_ser2(self, name, df_i, params):
        alp1 = Alpha_ser2(df_i)
        ary = eval(f'alp1.{name}')(params[0], params[1])
        return ary

    def eval_alpha_ser3(self, name, df_i, params):
        alp1 = Alpha_ser3(df_i)
        ary = eval(f'alp1.{name}')(params[0], params[1], params[2])
        return ary

    def eval_alpha_stat(self, name, df_i, params):
        alp1 = Alpha_ser_stastic(df_i)
        ary = eval(f'alp1.{name}')(params[0], params[1], params[2])
        return ary


class Df_main():
    
    def __init__(self, save_root=None) -> None:
        self.save_root = r'data/cmodty/my_feature/feature_data_df' if save_root is None else save_root     
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)

    def calc_alpha(self, func, name='alpha_ser_0', *params):
        try:
            t1 = time.time()
            df = func(name,  params)
            df = df[df.index>'2005-01-01']
            save_file = os.path.join(self.save_root, f'{name}_{params}.pkl.gzip')
            df.to_pickle(save_file)
            print(f"{name}_{params} calculating cost time:{(time.time()-t1):.2f}s")
        except:
            print("{name}_{params} calculating failed")

    def eval_alpha_df0(self, name, params):
        alp1 = Alpha_df0()
        ary = eval(f'alp1.{name}')()
        return ary

    def eval_alpha_df1(self, name, params):
        alp1 = Alpha_df1()
        ary = eval(f'alp1.{name}')(params[0])
        return ary

    def eval_alpha_df2(self, name, params):
        alp1 = Alpha_df2()
        ary = eval(f'alp1.{name}')(params[0], params[1])
        return ary

    def eval_alpha_df3(self, name, params):
        alp1 = Alpha_df3()
        ary = eval(f'alp1.{name}')(params[0], params[1], params[2])
        return ary

    def eval_alpha_df_stat(self, name, params):
        alp1 = Alpha_df_stat()
        ary = eval(f'alp1.{name}')(params[0], params[1])
        return ary

def calc_main_alpha_ser():
    #0: 15
    #1: 36
    #2: 17
    #3: 23
    #stat
    lst_v1 = list(range(39))
    lst_v2 = list(range(18))
    lst_v3 = list(range(12,24))
    param_lst = [2,5,10,30,60, 120,250]
    ratio_lst = [2,3,4]
    C = Ser_main()
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_ser0, f'alpha_ser0_{i}', 0) for i in range(16))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_ser1, f'alpha_ser1_{i}', j) for i,j in itertools.product(lst_v1, param_lst))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_ser2, f'alpha_ser2_{i}', j, k) for i,j, k in itertools.product(lst_v2, param_lst[:-2], ratio_lst))
    Parallel(n_jobs=16)(
        delayed(C.calc_alpha)( C.eval_alpha_stat, 'return_stastic', j, k, w )
        for j, k, w in itertools.product(
            ["max", "min", "mean", "median", "kurt", "skew"],
            param_lst,
            [ 'pepb_sz', 'pepb_sh', 'pepb_kc', 'pepb_cy', 'pepb_50', 'pepb_300', 'pepb_500', 'pepb_1000', ] ) )
    C3 = Ser_main(save_root=r'data/cmodty/my_feature/feature_data_ser3')
    Parallel(n_jobs=16)(delayed(C3.calc_alpha)(C3.eval_alpha_ser3, f'alpha_ser3_{i}', j, k, w) for i,j,k,w in itertools.product(lst_v3, param_lst[:-2], param_lst[:-2], param_lst[1:-2]))
    
    # for j, k, w in itertools.product(
    #         ["max", "min", "mean", "median", "kurt", "skew"],
    #         param_lst,
    #         [ 'pepb_sz', 'pepb_sh', 'pepb_kc', 'pepb_cy', 'pepb_50', 'pepb_300', 'pepb_500', 'pepb_1000', ] ):
    #         C.calc_alpha( C.eval_alpha_stat, 'return_stastic', j, k, w )
    
    # for i in range(12):
    #     C.calc_alpha(C.eval_alpha_ser0, f'alpha_ser0_{i}', 0)
    # for i,j in itertools.product(lst_v1, param_lst):
    #     C.calc_alpha(C.eval_alpha_ser1, f'alpha_ser1_{i}', j) 
    # for i,j, k in itertools.product(lst_v2, param_lst[:4], ratio_lst):
    #     C.calc_alpha(C.eval_alpha_ser2, f'alpha_ser2_{i}', j, k) 

def calc_main_alpha_df():
    #0: 13
    #1: 19
    #2: 7
    #3: 3
    #stat: 2
    lst_v1 = list(range(1,22))
    lst_v2 = list(range(8))
    lst_v3 = list(range(4))
    param_lst = [2,5,10,30,60, 120,250]
    ratio_lst = [2,3,4]
    C = Df_main()
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_df0, f'alpha_df0_{i}', 0) for i in range(12))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_df1, f'alpha_df1_{i}', j) for i,j in itertools.product(lst_v1, param_lst))
    Parallel(n_jobs=16)(delayed(C.calc_alpha)(C.eval_alpha_df2, f'alpha_df2_{i}', j, k) for i,j, k in itertools.product(lst_v2, param_lst[:-2], ratio_lst))
    Parallel(n_jobs=16)(
        delayed(C.calc_alpha)( C.eval_alpha_df_stat, f'alpha_df_stat_{i}',  k, w )
        for i, k, w in itertools.product(range(3), ["max", "min", "mean", "median", "kurt", "skew"], param_lst ) )
    CC = Df_main(r'data/cmodty/my_feature/feature_data_df3')
    Parallel(n_jobs=16)(delayed(CC.calc_alpha)(CC.eval_alpha_df3, f'alpha_df3_{i}', j, k, w) for i,j,k,w in itertools.product(lst_v3, param_lst[:-2], param_lst[:-2], param_lst[1:-2]))

    # for i in range(12):
    #     C.calc_alpha(C.eval_alpha_df0, f'alpha_df0_{i}', 0)
    # for i,j, k in itertools.product(lst_v2, param_lst[:4], ratio_lst):
    #     C.calc_alpha(C.eval_alpha_df2, f'alpha_df2_{i}', j, k) 


# calc_main_alpha_ser()
# calc_main_alpha_df()